import numpy as np
from typing import Dict, List, Tuple, Any

from rainforeLearn.gomoku.v2.configs.config import GomokuDQNConfig
from rainforeLearn.gomoku.v2.environment.threat_detector import ThreatDetector
from rainforeLearn.gomoku.v2.environment.game_rules import GameRules
from rainforeLearn.gomoku.v2.train.reward.reward_calculator import RewardCalculator


class GomokuEnvironment:
    """五子棋环境="""
    
    def __init__(self, config: GomokuDQNConfig):
        self.config = config
        self.board_size = config.board_size
        
        # 奖励配置
        self.invalid_move_penalty = config.invalid_move_penalty
        self.win_reward = config.win_reward
        self.lose_penalty = config.lose_penalty
        
        # 组件初始化
        self.threat_detector = ThreatDetector(self.board_size)
        self.reward_calculator = RewardCalculator(config)
        self.game_rules = GameRules(self.board_size)
        
        # 游戏状态
        self.reset()
    
    def reset(self) -> Dict[str, Any]:
        """重置游戏"""
        self.board = np.zeros((self.board_size, self.board_size), dtype=np.int8)
        self.current_player = 1  # 1为黑棋，-1为白棋
        self.game_over = False
        self.winner = 0
        self.move_count = 0
        return self.get_state()
    
    def get_state(self) -> Dict[str, Any]:
        """获取当前游戏状态"""
        board_view = self.board.copy()
        return {
            'board': board_view,
            'current_player': self.current_player,
            'legal_actions': self.game_rules.get_legal_actions(self.board),
            'board_flat': board_view.flatten().astype(np.float32)
        }
    
    def get_legal_actions(self) -> List[int]:
        """获取合法动作"""
        if self.game_over:
            return []
        return self.game_rules.get_legal_actions(self.board)
    
    def step(self, action: int) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
        """执行动作"""
        row, col = action // self.board_size, action % self.board_size
        
        # 检查动作合法性
        if not self.game_rules.is_valid_move(self.board, action):
            return self._handle_invalid_move()
        
        # 执行动作
        self.board[row, col] = self.current_player
        self.move_count += 1
        
        # 检查游戏结束
        reward, done, info = self._evaluate_game_state(row, col)
        
        # 切换玩家
        if not done:
            self.current_player = -self.current_player
        
        return self.get_state(), reward, done, info
    
    def _handle_invalid_move(self) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
        """处理非法动作"""
        self.game_over = True
        self.winner = -self.current_player
        
        info = {
            'winner': self.winner,
            'invalid_move': True,
            'original_reward': self.invalid_move_penalty,
            'move_count': self.move_count
        }
        
        return self.get_state(), self.invalid_move_penalty, True, info
    
    def _evaluate_game_state(self, row: int, col: int) -> Tuple[float, bool, Dict[str, Any]]:
        """评估游戏状态并计算奖励"""
        info = {
            'winner': None,
            'invalid_move': False,
            'move_count': self.move_count
        }
        
        # 检查获胜
        if self.game_rules.check_win(self.board, row, col, self.current_player):
            return self._handle_win(info)
        
        # 检查平局
        if self.game_rules.is_board_full(self.board):
            return self._handle_draw(info)
        
        # 计算中间奖励
        return self._calculate_intermediate_state(row, col, info)
    
    def _handle_win(self, info: Dict[str, Any]) -> Tuple[float, bool, Dict[str, Any]]:
        """处理获胜情况"""
        self.winner = self.current_player
        self.game_over = True
        info['winner'] = self.winner
        
        if self.current_player == 1:  # 黑棋获胜
            reward = self.win_reward
        else:  # 白棋获胜
            reward = self.lose_penalty
        
        info['original_reward'] = reward
        return reward, True, info
    
    def _handle_draw(self, info: Dict[str, Any]) -> Tuple[float, bool, Dict[str, Any]]:
        """处理平局情况"""
        self.winner = 0
        self.game_over = True
        info['winner'] = 0
        info['original_reward'] = 0.0
        return 0.0, False, info
    
    def _calculate_intermediate_state(self, row: int, col: int, 
                                    info: Dict[str, Any]) -> Tuple[float, bool, Dict[str, Any]]:
        """计算中间状态奖励"""
        if self.current_player == 1:  # 只对黑棋计算奖励
            # 威胁分析
            threat_info = self.threat_detector.analyze_threats(
                self.board, row, col, self.current_player
            )
            
            # 计算奖励
            reward = self.reward_calculator.calculate_intermediate_reward(
                self.board, row, col, self.current_player, 
                self.move_count, threat_info["threat_penalty"]
            )
            
            # 添加分析信息
            info.update(threat_info)
            info['attack_reward'] = self.reward_calculator._calculate_attack_reward(
                self.board, row, col, self.current_player
            )
            info['max_line'] = self.reward_calculator._get_max_line_length(
                self.board, row, col, self.current_player
            )
        else:
            reward = 0  # 白棋无奖励
        
        info['original_reward'] = reward
        return reward, False, info
